#include "map/occupancy_grid.h"

#include <ros/ros.h>

#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <algorithm>

#define CELL_VALUE_MAX 127
#define CELL_VALUE_MIN -127

#define PROBABILITY_MIN 0.0001f
#define PROBABILITY_MAX 0.9999f

#define LOGODDS_MAX 9.21024037
#define LOGODDS_SCALE ((float) (CELL_VALUE_MAX / LOGODDS_MAX))

void OccupancyGrid::writeToFile(std::ofstream& fileStream) {
	for (int x = 0; x < width_; x++) {
		for (int y = 0; y < height_; y++) {
			int n = grid_[x][y];
			fileStream << n << std::endl;
		}
	}
}

void OccupancyGrid::readFromFile(std::ifstream& fileStream) {
	for (int x = 0; x < width_; x++) {
		for (int y = 0; y < height_; y++) {
			int n;
			fileStream >> n;
			grid_[x][y] = n;
			if (n != 0) {
				dirty_rect_.lower_x = x < dirty_rect_.lower_x ? x : dirty_rect_.lower_x;
				dirty_rect_.lower_y = y < dirty_rect_.lower_y ? y : dirty_rect_.lower_y;
				dirty_rect_.upper_x = x > dirty_rect_.upper_x ? x : dirty_rect_.upper_x;
				dirty_rect_.upper_y = y > dirty_rect_.upper_y ? y : dirty_rect_.upper_y;
			}
		}
	}
}

OccupancyGrid::OccupancyGrid(int width, int height, float resolution) :
	resolution_(resolution),
	width_(width),
	height_(height)
{
	clearDirtyRect();

	// Build logodds table
	for (int i = 0; i < 256; i++) {
		logoddsLookup_[i] = logoddsToProbabilitySlow(i - 128);
	}

	// Tests
	int8_t logodds = probabilityToLogodds(0.5f);
	float prob = logoddsToProbability(logodds);
	if (fabs(0.5f - prob) > 0.01f) {
		ROS_ERROR("Occupancy grid test failed: converted prob 0.5 to logodds and back and got %f", prob);
	}

	// Allocate contigous memory for the occupancy grid and initialize it to 0
	grid_ = (int8_t **) malloc(sizeof(int8_t *) * width_);
	int8_t *tmp_grid = (int8_t *) malloc(sizeof(int8_t) * width_ * height_);
	memset(tmp_grid, probabilityToLogodds(0.5f), sizeof(int8_t) * width_ * height_);
	for (int i = 0; i < width_; i++) {
		grid_[i] = &tmp_grid[i * height_];
	}
}

OccupancyGrid::~OccupancyGrid()
{
	free(&grid_[0]);
	free(grid_);
}

void OccupancyGrid::clearDirtyRect() {
	dirty_rect_.lower_x = 10000;
	dirty_rect_.lower_y = 10000;
	dirty_rect_.upper_x = -10000;
	dirty_rect_.upper_y = -10000;
}

/**
 * Returns the probability from a cell's logodds representation.
 */
float OccupancyGrid::logoddsToProbabilitySlow(int8_t logodds) {
	float prob = 1.0f / (1.0f + expf(-logodds / LOGODDS_SCALE));
	return prob;
}

/**
 * Returns the probability from a cell's logodds representation.
 */
inline float OccupancyGrid::logoddsToProbability(int8_t logodds)
{
	return logoddsLookup_[128 + logodds];
}

/**
 * Returns a cell logodds representation for the given probability. In
 * the range [CELL_VALUE_MIN, CELL_VALUE_MAX]
 */
inline int8_t OccupancyGrid::probabilityToLogodds(float probability)
{
	if (probability < PROBABILITY_MIN) probability = PROBABILITY_MIN;
	if (probability > PROBABILITY_MAX) probability = PROBABILITY_MAX;

	float logodds = logf(probability) - logf(1.0f - probability);
	return LOGODDS_SCALE * logodds;
}

/**
 * Returns the probability of cell {cell_x, cell_y} is occupied.
 * The returned value is in the range [0.0, 1.0]
 */
float OccupancyGrid::probabilityOccupied(int cell_x, int cell_y)
{
	float prob = logoddsToProbability(grid_[cell_x][cell_y]);
	return prob;
}

/**
 * Updates the specified cell with the specified logodds value.
 */
void OccupancyGrid::update(int cell_x, int cell_y, int8_t value)
{
	// ROS_INFO("Updating cell (%d %d) with value %d", cell_x, cell_y, value);

	if (cell_x < 0 || cell_x >= width_ || cell_y < 0 || cell_y >= height_) return;

	dirty_rect_.lower_x = cell_x < dirty_rect_.lower_x ? cell_x : dirty_rect_.lower_x;
	dirty_rect_.lower_y = cell_y < dirty_rect_.lower_y ? cell_y : dirty_rect_.lower_y;
	dirty_rect_.upper_x = cell_x > dirty_rect_.upper_x ? cell_x : dirty_rect_.upper_x;
	dirty_rect_.upper_y = cell_y > dirty_rect_.upper_y ? cell_y : dirty_rect_.upper_y;

	int cell_value = grid_[cell_x][cell_y] + value;

	cell_value = std::min(CELL_VALUE_MAX, cell_value);
	cell_value = std::max(CELL_VALUE_MIN, cell_value);

	grid_[cell_x][cell_y] = cell_value;
}

/**
 * Returns the value of x on the normal distribution curve given the
 * specified mean and variance.
 */
static const float one_div_sqrtf_two_pi = 1.0f / sqrtf(2.0f * M_PI);
inline float normalDistribution(float x, float mean, float std_dev)
{
	float k = (x - mean) / std_dev;
	float probability_density = one_div_sqrtf_two_pi * expf(-0.5f * k * k);

	float factor = 1.0f / std_dev;
	return factor * probability_density;
}

/**
 * Adds a sensor ray into the grid. The ray originates at origin, goes in direction
 * and has the mean and standard deviation std_dev.
 */
void OccupancyGrid::addSensorRay(vec2 origin, float direction, float mean, float std_dev, float height)
{
	// Step size used for the marching
	float step = resolution_ / 2.0f;

	float cos_direction = cosf(direction);
	float sin_direction = sinf(direction);

	// Calculate the scale so that the height of the normal distribution curve is
	// 'scale' in height.
	float scale = height / normalDistribution(mean, mean, std_dev);

	// Calculate the end point for the ray. Ray continues three std_dev
	// after mean.
	float length = (mean + std_dev * 3.0f);

	// Step sizes to take along the ray
	float dx = step * cos_direction;
	float dy = step * sin_direction;

	float stepped_length = 0.0f;

	// Step along the ray and enter values read out from the normal
	// distribution into the cells
	while (stepped_length < length) {
		int cell_x = origin.x / resolution_;
		int cell_y = origin.y / resolution_;

		float probability = scale * normalDistribution(stepped_length, mean, std_dev);
		update(cell_x, cell_y, probabilityToLogodds(probability));

		origin.x += dx;
		origin.y += dy;

		stepped_length += step;
	}
}

/**
 * Returns the probability of this observed sensor ray. Value returned in the range [0, 1.0]
 */
float OccupancyGrid::probabilityOfSensorRay(vec2 origin, float direction, float mean, float std_dev, float height)
{
	// Step size used for the marching
	float step = resolution_ / 2.0f;

	float cos_direction = cosf(direction);
	float sin_direction = sinf(direction);

	// Calculate the scale so that the height of the normal distribution curve is
	// 'scale' in height.
	float scale = height / normalDistribution(mean, mean, std_dev);

	// Calculate the end point for the ray. Ray continues three std_dev
	// after mean.
	float length = (mean + std_dev * 3.0f);

	// Step sizes to take along the ray
	float dx = step * cos_direction;
	float dy = step * sin_direction;

	float stepped_length = 0.0f;

	// Step along the ray and measure difference in values read out from the normal
	// distribution compared to the cells
	float total_diff = 0.0f;
	int num_examined = 0;
	while (stepped_length < length) {
		int cell_x = origin.x / resolution_;
		int cell_y = origin.y / resolution_;

		float probability = scale * normalDistribution(stepped_length, mean, std_dev);
		total_diff += fabs(probability - logoddsToProbability(grid_[cell_x][cell_y]));

		origin.x += dx;
		origin.y += dy;

		stepped_length += step;
		num_examined++;
	}

	// A low diff should mean a higher probability for this sensor ray, so
	// invert, normalize and return
	return 1.0f - (total_diff / num_examined);
}

/**
 * Adds a symmetric 2D gaussian at the specified location.
 */
void OccupancyGrid::addHitWithDistribution(vec2 location, float std_dev, float height)
{
	//ROS_INFO("addHit with distribution at %.2f %.2f", location.x, location.y);

	// Calculate the scale so that the height of the normal distribution curve is
	// 'scale' in height.
	float scale = height / normalDistribution(0.0f, 0.0f, std_dev);

	int center_cell_x = location.x / resolution_;
	int center_cell_y = location.y / resolution_;

	// Go 1.5 standard deviations in each direction
	int cell_diff = ((1.5f * std_dev) / resolution_) + 0.5f;

	for (int x = center_cell_x - cell_diff; x < center_cell_x + cell_diff; x++) {
		for (int y = center_cell_y - cell_diff; y < center_cell_y + cell_diff; y++) {
			int x_dist = center_cell_x - x;
			int y_dist = center_cell_y - y;

			float dist = sqrtf(x_dist * x_dist + y_dist * y_dist);

			float probability = 0.5f + scale * normalDistribution(dist, 0, std_dev);
			update(x, y, probabilityToLogodds(probability));
		}
	}
}

/**
 * Subtracts probability across a ray. scale is the value to subtract in each cell.
 */
void OccupancyGrid::subtractRay(vec2 from, vec2 to, float scale)
{

}

/**
 * Adds a line 'beam' to the occupancy grid. It starts at 'from', goes to 'to'.
 */
void OccupancyGrid::add(vec2 from, vec2 to, bool hit)
{
	// TODO: This is just a simple line tracer without
	// an inverse sensor model. This will need an inverse sensor
	// model later on.

	float step = 0.01f;

	int end_cell_x = to.x / resolution_;
	int end_cell_y = to.y / resolution_;

	float curr_x = from.x;
	float curr_y = from.y;

	int curr_cell_x = from.x / resolution_;
	int curr_cell_y = from.y / resolution_;

	int8_t occupiedLogodds = probabilityToLogodds(1.0f);
	int8_t emptyLogodds = probabilityToLogodds(0.0f);

	float t = 0.0f;
	while (t <= 1.0f) {
		curr_x = t * to.x + (1.0f - t) * from.x;
		curr_y = t * to.y + (1.0f - t) * from.y;

		curr_cell_x = curr_x / resolution_;
		curr_cell_y = curr_y / resolution_;

		if (curr_cell_x == end_cell_x && curr_cell_y == end_cell_y) {
			update(curr_cell_x, curr_cell_y, occupiedLogodds);
		} else {
			update(curr_cell_x, curr_cell_y, emptyLogodds);
		}

		t += step;
	}
}
